#ifndef TRACKER_H
#define TRACKER_H

#include <vector>
#include <tuple>
#include <numeric>
#include "kalman_tracker.h"
#include "hungarian.h"
#include "utils.h"



using DistanceMetricFunc = std::function<
        vector<vector<float>>(const std::vector<int> &trk_ids, const std::vector<int> &det_ids)>;

const float INVALID_DIST = 1E3f;

void associate_detections_to_trackers_idx(const DistanceMetricFunc &metric,
                                          std::vector<int> &unmatched_trks,
                                          std::vector<int> &unmatched_dets,
                                          std::vector<std::tuple<int, int>> &matched);

template<typename TrackData>
class TrackerManager {
public:
    explicit TrackerManager(std::vector<TrackData> &data) : data(data) {}

    void predict() {
        for (auto &t:data) {
            t.kalman.predict();
        }
    }

    void remove_nan() {
        data.erase(remove_if(data.begin(), data.end(),
                             [](const TrackData &t) {
                                 auto bbox = t.kalman.rect();
                                 return std::isnan(bbox.x) || std::isnan(bbox.y) ||
                                        std::isnan(bbox.width) || std::isnan(bbox.height);
                             }),
                   data.end());
    }

    void remove_deleted() {
        data.erase(remove_if(data.begin(), data.end(),
                             [this](const TrackData &t) {
                                 return t.kalman.state() == TrackState::Deleted;
                             }), data.end());
    }

    std::vector<std::tuple<int, int>>
    update(const std::vector<BBox> &dets,
           const DistanceMetricFunc &confirmed_metric, const DistanceMetricFunc &unconfirmed_metric) {
        std::vector<int> unmatched_trks;
        for (size_t i = 0; i < data.size(); ++i) {
            if (data[i].kalman.state() == TrackState::Confirmed || 
                data[i].kalman.state() == TrackState::Covered) {
                unmatched_trks.emplace_back(i);
            }
        }

        std::vector<int> unmatched_dets(dets.size());
        iota(unmatched_dets.begin(), unmatched_dets.end(), 0);

        std::vector<std::tuple<int, int>> matched;

        associate_detections_to_trackers_idx(confirmed_metric, unmatched_trks, unmatched_dets, matched);

        for (size_t i = 0; i < data.size(); ++i) {
            if (data[i].kalman.state() == TrackState::Tentative) {
                unmatched_trks.emplace_back(i);
            }
        }

        associate_detections_to_trackers_idx(unconfirmed_metric, unmatched_trks, unmatched_dets, matched);

        for (auto i : unmatched_trks) {
            data[i].kalman.miss();
        }

        // update matched trackers with assigned detections.
        // each prediction is corresponding to a manager
        for (auto[x, y] : matched) {
            data[x].kalman.update(dets[y].rect);
            data[x].kalman.setClsId(dets[y].cls_id);
        }

        // create and initialise new trackers for unmatched detections
        for (auto umd : unmatched_dets) {
            matched.emplace_back(data.size(), umd);
            auto t = TrackData{};
            t.kalman.init(dets[umd].rect);
            data.emplace_back(t);
        }

        return matched;
    }

    void visible_tracks(std::vector<BBox>& trackBoxes) {
        trackBoxes = std::vector<BBox>();
        for (size_t i = 0; i<data.size(); ++i) {
            auto& t = data[i];   // 注意，这里要用&，否则构造一个deepsort中使用的TrackData对象将会消耗0.2ms
            auto bbox = t.kalman.rect();
            if (bbox.x < img_box.x) bbox.x = img_box.x;
            if (bbox.y < img_box.y) bbox.y = img_box.y;
            if (bbox.width + bbox.x > img_box.width) bbox.width = img_box.width - bbox.x;
            if (bbox.height + bbox.y > img_box.height) bbox.height = img_box.height - bbox.y;
            if (t.kalman.state() == TrackState::Confirmed) {
                BBox res;
                res.trk_id = t.kalman.getTrkId();
                res.cls_id = t.kalman.getClsId();
                res.rect = bbox;
                trackBoxes.push_back(res);
            }
        }
    }

    void Set(uint32_t width, uint32_t height) {
        img_box.x = 0;
        img_box.y = 0;
        img_box.width = width;
        img_box.height = height;
    }

private:
    std::vector<TrackData> &data;
    cv::Rect img_box;
};

#endif //TRACKER_H
